import os

# llm
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate

# retriever
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools.retriever import create_retriever_tool

# web search
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.utilities import BingSearchAPIWrapper

# agent
from langchain.agents import AgentExecutor, create_tool_calling_agent

from .utils import read_template, parse_rules


class SuricataToolAgent:
    def __init__(self, agent_executor: AgentExecutor):
        self.agent_executor = agent_executor

    @classmethod
    def construct(cls, writing_guide_path: str, prompt_template_path: str, llm: ChatOpenAI):
        # Build a retriever as a tool
        loader = UnstructuredMarkdownLoader(writing_guide_path)
        suricata_rules = loader.load()

        all_splits = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200).split_documents(
            suricata_rules)
        vectorstore = Chroma.from_documents(documents=all_splits, embedding=OpenAIEmbeddings())
        retriever = vectorstore.as_retriever(search_type="mmr")

        retriever_tool = create_retriever_tool(
            retriever,
            "Suricata_Rules_Writing_Guide",
            "Search for writing guide about Suricata Rules. When you need to write Suricata Rules, you must use this "
            "tool to"
            "obtain the writing rules."
        )

        # Build web search tool
        des = (
            "A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to "
            "answer"
            "question about writing Suricata Rules. Input should be a search query.")
        search = TavilySearchResults(max_results=10, description=des)

        # Create the Agent
        tools = [search, retriever_tool]
        template = read_template(prompt_template_path)
        prompt = PromptTemplate.from_template(template)

        agent = create_tool_calling_agent(llm, tools, prompt)
        agent_executor = AgentExecutor(agent=agent, tools=tools)
        return cls(agent_executor)

    def invoke(self, suricata_des: str) -> list[str]:
        llm_response = self.agent_executor.invoke({"input": suricata_des})['output']
        return parse_rules(llm_response)


os.environ["http_proxy"] = 'http://127.0.0.1:7890'
os.environ["https_proxy"] = 'http://127.0.0.1:7890'
os.environ["OPENAI_API_BASE"] = 'https://api.xty.app/v1'
os.environ["OPENAI_API_KEY"] = 'sk-8GXQC4rR3kslxPD5398bF0D9B46c4bE89f76537dC36aE1Cb'
os.environ['TAVILY_API_KEY'] = "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"

llm = ChatOpenAI(model='gpt-4')
tool_agent = SuricataToolAgent.construct('./writing_guide/suricata_rules.md', './prompt/tool_agent_template.txt', llm)
